local Class = require('lua.lang.Class')
local Edge = require('lua.cfg.Edge')
local StringBuilder = require('lua.lang.StringBuilder')
local EdgeType = require('lua.cfg.EdgeType')
local Lua = require('lua.vm.Lua')

--- @class ControlFlowGraph : Class
--- @field startBlock BasicBlock
--- @field blocks table<number, BasicBlock>
--- @field edges table<Edge>
--- @field labels table<number, string>
--- @field blockToIndex table<BasicBlock, number>
--- @field nextBlockId number
local ControlFlowGraph = Class {}

--- @return ControlFlowGraph
function ControlFlowGraph:new()
	--- @type ControlFlowGraph
	local o = {}
	o.blocks = {}
	o.edges = {}
	o.labels = {}
	o.blockToIndex = {}
	o.nextBlockId = 0
	return setmetatable(o, ControlFlowGraph)
end

--- @param source BasicBlock
--- @return number
function ControlFlowGraph:getInDegree(source)
	local edges = self.edges
	local count = 0
	for _, edge in pairs(edges) do
		if edge.target == source then
			count = count + 1
		end
	end
	return count
end

--- @param target BasicBlock
--- @return number
function ControlFlowGraph:getOutDegree(target)
	local edges = self.edges
	local count = 0
	for _, edge in pairs(edges) do
		if edge.source == target then
			count = count + 1
		end
	end
	return count
end

--- @param source BasicBlock
--- @return Edge[]
function ControlFlowGraph:getIncomingEdges(source)
	local edges = self.edges
	--- @type Edge[]
	local set = {}
	for _, edge in pairs(edges) do
		if edge.target == source then
			set[#set + 1] = edge
		end
	end
	return set
end

--- @param target BasicBlock
--- @return Edge[]
function ControlFlowGraph:getOutgoingEdges(target)
	local edges = self.edges
	--- @type Edge[]
	local set = {}
	for _, edge in pairs(edges) do
		if edge.source == target then
			set[#set + 1] = edge
		end
	end
	return set
end

--- @param block BasicBlock
--- @param label string?
--- @return void
function ControlFlowGraph:addBlock(block, label)
	local i = self.nextBlockId + 1
	label = label or 'B' .. i
	self.nextBlockId = i
	self.blocks[i] = block
	self.labels[i] = label
	self.blockToIndex[block] = i
end

--- @param source BasicBlock
--- @param target BasicBlock
--- @param type EdgeType
--- @return void
function ControlFlowGraph:addEdge(source, target, type)
	local edges = self.edges
	edges[#edges + 1] = Edge:new(source, target, type)
end

--- @param source BasicBlock
--- @param target BasicBlock
--- @return void
function ControlFlowGraph:removeEdge(source, target)
	local edges = self.edges
	for i = #edges, 1, -1 do
		local edge = edges[i]
		if edge.source == source and edge.target == target then
			--- @type Edge
			local e = table.remove(edges, i)
			assert(e.source == source and e.target == target)
			return e
		end
	end
end

--- @return string
function ControlFlowGraph:toDotString()
	local digraph = StringBuilder:new()
	digraph:append('digraph ControlFlowGraph {\n')
		   :append('\tlabel=\34ControlFlowGraph\34;\n')
		   :append('\tlabelloc=t;\n')
		   :append('\trankdir=TB;\n')
		   :append('\tnode [shape=rectangle, style=rounded];\n\n')
	local labels = self.labels
	local blockToIndex = self.blockToIndex
	for idx, block in pairs(self.blocks) do
		local label = labels[idx]
		local sb = StringBuilder:new()
		local instructions = block.instructions
		for i = 1, #instructions do
			sb:append(instructions[i]:toString()):append('\\n')
		end
		digraph:append('\t' .. label .. ' [label=\34' .. sb:toString() .. '\34];\n')
	end
	digraph:append('\tstart [shape=point];\n')
			:append('\tstart -> ' .. labels[blockToIndex[self.startBlock]] .. ';\n\n')
	for _, edge in pairs(self.edges) do
		local sourceLabel = labels[blockToIndex[edge.source]]
		local targetLabel = labels[blockToIndex[edge.target]]
		local type = edge.type
		local edgeAttributes
		if type == EdgeType.NEXT then
			edgeAttributes = '[color=black]'
		elseif type == EdgeType.JUMP then
			edgeAttributes = '[color=red, style=dashed]'
		else
			edgeAttributes = '[color=black]'
		end
		digraph:append('\t' .. sourceLabel .. ' -> ' .. targetLabel .. ' ' .. edgeAttributes .. ';\n')
	end
	return digraph:append('}\n'):toString()
end

--- @param source BasicBlock
--- @param target BasicBlock
--- @return void
function ControlFlowGraph:merge(source, target)
	local sourceInstructions = source.instructions
	local targetInstructions = target.instructions
	sourceInstructions[#sourceInstructions] = nil
	for i = 1, #targetInstructions do
		sourceInstructions[#sourceInstructions + 1] = targetInstructions[i]
	end
	local outgoingEdges = self:getOutgoingEdges(target)
	self:removeEdge(source, target)
	local blocks = self.blocks
	local blockToIndex = self.blockToIndex
	blocks[blockToIndex[target]] = nil
	blockToIndex[target] = nil
	for i = 1, #outgoingEdges do
		local edge = outgoingEdges[i]
		self:removeEdge(edge.source, edge.target)
	end
	for i = 1, #outgoingEdges do
		local edge = outgoingEdges[i]
		self:addEdge(source, edge.target, edge.type)
	end
end

--- @return void
function ControlFlowGraph:optimize()
	local blocks = self.blocks
	--- @type BasicBlock[]
	for _, block in pairs(blocks) do
		local exit = block:getExitInstruction()
		if exit == nil then
			goto forContinue
		end
		local canNext = exit:getOpCode() == Lua.OP_JMP and exit:getArgA() > 0
		while not canNext do
			local outgoingEdges = self:getOutgoingEdges(block)
			if #outgoingEdges == 1 then
				local target = outgoingEdges[1].target
				if self:getInDegree(target) == 1 then
					self:merge(block, target)
					goto whileContinue
				end
			end
			canNext = true
			:: whileContinue ::
		end
		:: forContinue ::
	end
end

return ControlFlowGraph
